"""
数据获取模块
负责从akshare获取茅台股票数据并进行预处理
"""

import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import streamlit as st


# 常用股票配置
POPULAR_STOCKS = {
    "贵州茅台": "600519",
    "中国平安": "601318", 
    "招商银行": "600036",
    "五粮液": "000858",
    "腾讯控股": "700.HK",  # 港股需要特殊处理
    "比亚迪": "002594",
    "宁德时代": "300750",
    "中国移动": "600941",
    "阿里巴巴-SW": "9988.HK",  # 港股
    "美的集团": "000333",
    "工商银行": "601398",
    "建设银行": "601939",
    "中国石化": "600028",
    "中国石油": "601857",
    "万科A": "000002",
    "格力电器": "000651",
    "海康威视": "002415",
    "立讯精密": "002475",
    "药明康德": "603259",
    "隆基绿能": "601012"
}

# 行业分类
INDUSTRY_STOCKS = {
    "白酒类": {
        "贵州茅台": "600519",
        "五粮液": "000858",
        "剑南春": "600779",
        "泸州老窖": "000568",
        "山西汾酒": "600809"
    },
    "银行类": {
        "招商银行": "600036",
        "工商银行": "601398",
        "建设银行": "601939",
        "农业银行": "601288",
        "中国银行": "601988"
    },
    "科技类": {
        "宁德时代": "300750",
        "比亚迪": "002594",
        "海康威视": "002415",
        "立讯精密": "002475",
        "韦尔股份": "603501"
    },
    "地产类": {
        "万科A": "000002",
        "保利发展": "600048",
        "中国恒大": "600340",
        "融创中国": "600918",
        "碧桂园": "002007"
    },
    "消费类": {
        "格力电器": "000651",
        "美的集团": "000333",
        "海尔智家": "600690",
        "伊利股份": "600887",
        "中国中免": "601888"
    }
}

def get_stock_options():
    """
    获取股票选择选项
    
    Returns:
        tuple: (股票名称列表, 股票代码字典)
    """
    return list(POPULAR_STOCKS.keys()), POPULAR_STOCKS


def get_industry_options():
    """
    获取行业分类选项
    
    Returns:
        dict: 行业分类字典
    """
    return INDUSTRY_STOCKS


def validate_stock_code(stock_code):
    """
    验证股票代码格式并获取股票信息
    
    Args:
        stock_code: 股票代码
        
    Returns:
        tuple: (是否有效, 股票信息字典)
    """
    import re
    
    stock_code = stock_code.strip().upper()
    
    # A股代码格式: 6位数字
    if re.match(r'^\d{6}$', stock_code):
        try:
            # 使用akshare获取股票基本信息
            # 方法1: 尝试获取股票信息
            try:
                stock_info_df = ak.stock_individual_info_em(symbol=stock_code)
                if len(stock_info_df) > 0:
                    # 查找股票名称
                    stock_name = "未知股票"
                    for _, row in stock_info_df.iterrows():
                        if '股票简称' in str(row.get('item', '')) or '名称' in str(row.get('item', '')):
                            stock_name = str(row.get('value', '未知股票'))
                            break
                    
                    return True, {
                        'code': stock_code,
                        'name': stock_name,
                        'type': 'A股',
                        'market': '沪深'
                    }
            except:
                pass
            
            # 方法2: 使用股票代码前缀判断市场并设置默认名称
            market_info = get_market_info(stock_code)
            if market_info:
                # 尝试通过实时数据验证股票存在
                try:
                    # 获取少量历史数据来验证股票代码有效性
                    from datetime import datetime, timedelta
                    end_date = datetime.now().strftime('%Y%m%d')
                    start_date = (datetime.now() - timedelta(days=7)).strftime('%Y%m%d')
                    
                    test_df = ak.stock_zh_a_hist(
                        symbol=stock_code, 
                        period="daily", 
                        start_date=start_date, 
                        end_date=end_date, 
                        adjust=""
                    )
                    
                    if len(test_df) > 0:
                        return True, {
                            'code': stock_code,
                            'name': f"{market_info['market']}{stock_code}",
                            'type': 'A股',
                            'market': market_info['market']
                        }
                except:
                    pass
                    
        except Exception as e:
            st.warning(f"获取股票信息时出错: {e}")
    
    # 港股代码格式: 数字.HK
    if re.match(r'^\d+\.HK$', stock_code):
        return True, {
            'code': stock_code,
            'name': f'港股{stock_code}',
            'type': '港股',
            'market': '香港'
        }
    
    return False, {}


def get_market_info(stock_code):
    """
    根据股票代码获取市场信息
    
    Args:
        stock_code: 6位股票代码
        
    Returns:
        dict: 市场信息
    """
    code_prefix = stock_code[:3]
    
    # 沪市主板
    if code_prefix in ['600', '601', '602', '603', '605']:
        return {'market': '沪市', 'board': '主板'}
    
    # 深市主板
    elif code_prefix in ['000', '001', '002']:
        return {'market': '深市', 'board': '主板'}
    
    # 创业板
    elif code_prefix in ['300', '301']:
        return {'market': '创业板', 'board': '创业板'}
    
    # 科创板
    elif code_prefix in ['688', '689']:
        return {'market': '科创板', 'board': '科创板'}
    
    # 北交所
    elif code_prefix in ['430', '831', '832', '833', '834', '835', '836', '837', '838', '839']:
        return {'market': '北交所', 'board': '北交所'}
    
    return None


def get_stock_name_by_code(stock_code):
    """
    通过多种方式获取股票名称
    
    Args:
        stock_code: 股票代码
        
    Returns:
        str: 股票名称
    """
    try:
        # 方法1: 从预定义的股票列表中查找
        for name, code in POPULAR_STOCKS.items():
            if code == stock_code:
                return name
        
        # 方法2: 从行业分类中查找
        for industry, stocks in INDUSTRY_STOCKS.items():
            for name, code in stocks.items():
                if code == stock_code:
                    return name
        
        # 方法3: 使用akshare的股票信息接口
        try:
            stock_info = ak.stock_individual_info_em(symbol=stock_code)
            for _, row in stock_info.iterrows():
                item = str(row.get('item', ''))
                if '简称' in item or '名称' in item:
                    return str(row.get('value', f'股票{stock_code}'))
        except:
            pass
            
        # 方法4: 返回市场信息 + 代码
        market_info = get_market_info(stock_code)
        if market_info:
            return f"{market_info['market']}股票{stock_code}"
            
    except:
        pass
    
    return f"股票{stock_code}"


def search_stock_by_keyword(keyword):
    """
    根据关键词搜索股票
    
    Args:
        keyword: 搜索关键词
        
    Returns:
        list: 匹配的股票列表
    """
    try:
        # 使用akshare搜索股票
        search_results = ak.tool_trade_date_hist_sina()
        # 这里可以添加更复杂的搜索逻辑
        return []
    except:
        return []


def get_stock_suggestions(input_text):
    """
    获取股票代码建议
    
    Args:
        input_text: 用户输入文本
        
    Returns:
        list: 建议的股票代码列表
    """
    suggestions = []
    
    # 从热门股票中匹配
    for name, code in POPULAR_STOCKS.items():
        if input_text in name or input_text in code:
            suggestions.append(f"{name}({code})")
    
    # 从行业股票中匹配
    for industry, stocks in INDUSTRY_STOCKS.items():
        for name, code in stocks.items():
            if input_text in name or input_text in code:
                suggestions.append(f"{name}({code}) - {industry}")
    
    return suggestions[:10]  # 限制返回10个建议



    """计算RSI相对强弱指数"""
    delta = prices.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi


def calculate_rsi(prices, window=14):
    """计算RSI相对强弱指数"""
    delta = prices.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))
    return rsi


@st.cache_data(ttl=3600)  # 缓存1小时
def get_stock_data(stock_code="600519", stock_name="茅台", days=730):
    """
    获取股票数据
    
    Args:
        stock_code: 股票代码（不含后缀）
        stock_name: 股票名称
        days: 获取过去多少天的数据，默认730天(2年)
    
    Returns:
        tuple: (DataFrame, 数据类型)
    """
    try:
        # 计算日期范围
        end_date = datetime.now().strftime('%Y%m%d')
        start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
        
        # 获取股票数据
        df = ak.stock_zh_a_hist(
            symbol=stock_code, 
            period="daily", 
            start_date=start_date, 
            end_date=end_date, 
            adjust="qfq"
        )
        
        if len(df) == 0:
            raise Exception("获取到空数据")
        
        # 重命名列
        df = df.rename(columns={
            '日期': 'date',
            '开盘': 'open',
            '收盘': 'close', 
            '最高': 'high',
            '最低': 'low',
            '成交量': 'volume'
        })
        
        # 转换日期格式
        df['date'] = pd.to_datetime(df['date'])
        df = df.set_index('date')
        
        return df, "真实数据"
        
    except Exception as e:
        st.warning(f"获取{stock_name}({stock_code})真实数据失败: {e}，使用模拟数据")
        return create_mock_data(stock_name), "模拟数据"


def create_mock_data(stock_name="茅台"):
    """创建模拟股票数据"""
    np.random.seed(42)
    
    # 生成时间序列
    end_date = datetime.now()
    start_date = end_date - timedelta(days=730)
    dates = pd.date_range(start=start_date, end=end_date, freq='D')
    dates = [d for d in dates if d.weekday() < 5]  # 只保留工作日
    
    base_price = 1650
    n_days = len(dates)
    
    # 生成随机价格
    price_changes = np.random.normal(0, 20, n_days)
    prices = [base_price]
    
    for change in price_changes[1:]:
        new_price = prices[-1] + change
        new_price = max(1200, min(2200, new_price))
        prices.append(new_price)
    
    # 生成OHLC数据
    opens, highs, lows, closes, volumes = [], [], [], prices, []
    
    for i, close in enumerate(closes):
        if i == 0:
            open_price = close + np.random.normal(0, 5)
        else:
            open_price = closes[i-1] + np.random.normal(0, 10)
        
        daily_range = abs(np.random.normal(0, 30))
        high = max(open_price, close) + daily_range * 0.7
        low = min(open_price, close) - daily_range * 0.3
        volume = max(100000, int(np.random.normal(500000, 200000)))
        
        opens.append(open_price)
        highs.append(high)
        lows.append(low)
        volumes.append(volume)
    
    df = pd.DataFrame({
        'open': opens,
        'high': highs,
        'low': lows,
        'close': closes,
        'volume': volumes
    }, index=dates)
    
    return df


def add_technical_indicators(df):
    """
    添加技术指标
    
    Args:
        df: 股票数据DataFrame
    
    Returns:
        DataFrame: 添加技术指标后的数据
    """
    df = df.copy()
    
    # 移动平均线
    df['ma5'] = df['close'].rolling(window=5).mean()
    df['ma20'] = df['close'].rolling(window=20).mean()
    df['ma60'] = df['close'].rolling(window=60).mean()
    
    # RSI指标
    df['rsi'] = calculate_rsi(df['close'])
    
    # 价格变化
    df['previous_close'] = df['close'].shift(1)
    df['price_change'] = df['close'] - df['previous_close']
    df['price_change_pct'] = (df['price_change'] / df['previous_close']) * 100
    
    # 成交量移动平均
    df['volume_ma'] = df['volume'].rolling(window=5).mean()
    
    # MACD指标
    exp1 = df['close'].ewm(span=12).mean()
    exp2 = df['close'].ewm(span=26).mean()
    df['macd'] = exp1 - exp2
    df['macd_signal'] = df['macd'].ewm(span=9).mean()
    df['macd_histogram'] = df['macd'] - df['macd_signal']
    
    # 布林带
    rolling_mean = df['close'].rolling(window=20).mean()
    rolling_std = df['close'].rolling(window=20).std()
    df['bollinger_upper'] = rolling_mean + (rolling_std * 2)
    df['bollinger_lower'] = rolling_mean - (rolling_std * 2)
    
    # 删除NaN值
    df = df.dropna()
    
    return df


def get_stock_info(df, stock_name="股票"):
    """
    获取股票基本信息
    
    Args:
        df: 股票数据DataFrame
    
    Returns:
        dict: 股票基本信息
    """
    if len(df) == 0:
        return {}
    
    current_price = df['close'].iloc[-1]
    previous_close = df['close'].iloc[-2] if len(df) > 1 else current_price
    
    info = {
        'current_price': current_price,
        'previous_close': previous_close,
        'change': current_price - previous_close,
        'change_pct': ((current_price - previous_close) / previous_close) * 100,
        'high_52w': df['high'].max(),
        'low_52w': df['low'].min(),
        'volume_avg': df['volume'].mean(),
        'data_start': df.index[0].strftime('%Y-%m-%d'),
        'data_end': df.index[-1].strftime('%Y-%m-%d'),
        'total_records': len(df)
    }
    
    # 技术指标当前值
    if 'ma5' in df.columns:
        info['ma5'] = df['ma5'].iloc[-1]
    if 'ma20' in df.columns:
        info['ma20'] = df['ma20'].iloc[-1]
    if 'rsi' in df.columns:
        info['rsi'] = df['rsi'].iloc[-1]
        
    return info